#ifndef CLASS_HALO
#define CLASS_HALO

#include <stdio.h>
#include <math.h>

#include "P.h"
#include "AG.h"
#include "TreeForceFinder.h"
#include "Vector.h"


#define SQR(x) ((x)*(x))

typedef unsigned int uint;
typedef long unsigned int luint;

using namespace std;

	
	class Halo { // ( Hernquist )
	public:
		
		uint n; 
		double m; // total mass
		double r0; // scale length of Hernquist profile 
		double rmax;
		
		double G;

		P *p;
		
		Halo():
		G(1),
		p(0)
		{ }
		
		
		virtual void set_pos(P *p, int n) {
			
			printf("Halo::set_pos()... "); fflush(stdout);
			
			srand48(222);
			
			double r;
			
			for (uint i=0; i<n; i++) {
				
				do {
					double q = drand48(); 
					r = r0 * (q + sqrt(q)) / (1 - q);
				} while (r > rmax);
				
				double cosa1 = 1-2*drand48();
				double a1 = acos(cosa1); 
				double a2 = 2*M_PI*drand48(); 
				
				p[i].Pos[0] = r*sin(a1)*cos(a2);
				p[i].Pos[1] = r*sin(a1)*sin(a2);
				p[i].Pos[2] = r*cosa1;
				
				p[i].pot = pot(r);
			}
			
			for(int i=0; i<n; i++)
				p[i].Mass = m/n;

			printf("done\n");
		}
		
		virtual void set_pot(P *p, int n) {
			
			for (int i=0; i<n; i++) {
				p[i].pot = pot( SQR(p[i].Pos[0])+SQR(p[i].Pos[1])+SQR(p[i].Pos[2]) );	
			}
			
		}
		
		
		double rho(double r) {
			return m/(2*M_PI)*r0/(r*pow(r+r0, 3));
		};
		
		double M(double r1, double z1, double r2, double z2, int nr=1, int nz=1) {
			
			double m(0);
			
			double dr = (r2-r1)/nr;
			double dz = (z2-z1)/nz;
			
			for (int iz=0; iz< nz; iz++)
			for (int ir=0; ir< nr; ir++) {
				
				double rci = r1 + dr*(ir+0.5);
				double zci = z1 + dz*(iz+0.5);
				
				double rzi = sqrt(SQR(rci) + SQR(zci));
				
				m+=rho(rzi)*rci;
				
			}
			
			return 4*M_PI*m*dr*dz;
			
		}
		
		double M(double r1, double r2,  int nr=1) {
			
			double m(0);
			double dr = (r2-r1)/nr;
			
			for (int ir=0; ir< nr; ir++) {
				
				double r1i = r1 + dr*(ir+0.0);
				double rci = r1 + dr*(ir+0.5);
				double r2i = r1 + dr*(ir+1.0);
				
				m+= rho(rci)*(pow(r2i,3)-pow(r1i, 3));
						
				//printf("%g %g\n", rci, halo.rho(rci));
				
			}
			
			return 4.0/3*M_PI*m;
		}

		
		
		
		inline double acc(double r) {
			
			double rr0 = r+r0;
			return -G*m/(rr0*rr0);
			
		}
		
		
		double pot(double r){
			return -G*m/(r+r0);			
		}
		
		
		double PDF(double e) {
			
			double q = sqrt(-r0/(G*m)*e);
			double vg = sqrt(G*m/r0);
			return m / ( 8*sqrt(2)*pow(M_PI*r0*vg,3) ) * 1.0/pow(1-q*q, 2.5)*( 3*asin(q) + 
								q*sqrt(1-q*q)*(1-2*q*q)*(8*q*q*q*q - 8*q*q - 3) );
			
		};
		
		double g(double r, double v){
			
			double Epj = -G*m/(r+r0);
			double Ekj = .5*v*v;
			double Ej  = Epj + Ekj;
			if (Ej>0) Ej=0; 
					
			return PDF(Ej)/rho(r);
			
		}
		
		double vesc(double r) {
			return sqrt(2*G*m/(r+r0));
		}
		
		template <typename T> 
		void set_vel(T *p, int n) {
		
			printf("Halo::set_vel()... "); fflush(stdout);
			
			int i(0);
			while ( i<n) {

				double r = sqrt(SQR(p[i].Pos[0]) + SQR(p[i].Pos[1]) + SQR(p[i].Pos[2]));
				
				double v_esc = vesc(r);
				p[i].Vel[0] = v_esc*(2*drand48() - 1);
				p[i].Vel[1] = v_esc*(2*drand48() - 1);
				p[i].Vel[2] = v_esc*(2*drand48() - 1);
				
				double v = sqrt(SQR(p[i].Vel[0]) + SQR(p[i].Vel[1]) + SQR(p[i].Vel[2]));
				if ( v_esc < v ) continue; 

				double gi =  g(r,0)*drand48();
				if ( g(r,v) < gi  ) continue; 
				
				i++;
		
			}
			
			printf("done\n");
 		}

		void set_vel_uniform(P *p, int n) {
			
			printf("Halo::set_vel_uniform()... "); fflush(stdout);
			
			int i(0);
			while ( i<n) {
				
				double r = sqrt(SQR(p[i].Pos[0]) + SQR(p[i].Pos[1]) + SQR(p[i].Pos[2]));
				double v_esc = vesc(r);
				
				p[i].Vel[0] = v_esc*(2*drand48() - 1);
				p[i].Vel[1] = v_esc*(2*drand48() - 1);
				p[i].Vel[2] = v_esc*(2*drand48() - 1);
				
				double v = sqrt(SQR(p[i].Vel[0]) + SQR(p[i].Vel[1]) + SQR(p[i].Vel[2]));
				if ( v_esc < v ) continue; 
				
				i++;
			}
			printf("done\n");
		}
		
		
		void set_vesc(P *p, int n, AG &ag) {

			printf("Halo::set_vesc(p, %i, ag)...\n", n); fflush(stdout);
			
			int i(0);
			for (int i=0; i<n; i++){

				double r = sqrt( SQR(p[i].Pos[0]) + SQR(p[i].Pos[1]) );
				double z = fabs(p[i].Pos[2]);
				
				AG::A a = ag.get_acc_pot(r,z);
				p[i].v_esc = sqrt(-2*a.pot);
			
			}
			
			printf("done\n");
			
		}
				
		void set_vel_uniform(P *p, int n, AG &ag) {

			printf("Halo::set_vel_uniform(p, %i, ag)...\n", n); fflush(stdout);
			
			int i(0);
			while ( i<n ) {

				double r = sqrt( SQR(p[i].Pos[0]) + SQR(p[i].Pos[1]) );
				double z = fabs(p[i].Pos[2]);
				
				AG::A a = ag.get_acc_pot(r,z);
				double v_esc = sqrt(-2*a.pot);
				
				p[i].Vel[0] = v_esc*(2*drand48() - 1);
				p[i].Vel[1] = v_esc*(2*drand48() - 1);
				p[i].Vel[2] = v_esc*(2*drand48() - 1);
				p[i].v_esc = v_esc;
				p[i].pot = a.pot;
				//printf("%g\n", v_esc);
				
				double v = sqrt(SQR(p[i].Vel[0]) + SQR(p[i].Vel[1]) + SQR(p[i].Vel[2]));
				if ( v_esc < v ) continue; 
				
				i++;
			}
			
			printf("done\n");
			
		}
		
		
		void calc_and_save_w(	const char *fnr ) {
			
			printf("Halo::calc_and_save_w(%s)...", fnr); fflush(stdout);

			
			class SG { public:
				uint nr;
				double r1, r2;
				SG(): nr(8),r1(1),r2(8192) {};
			}sg;
			
			class VG { public:
				uint nr, nt;
				VG(): nr(20),nt(20) {};
			}vg;
			
			uint wnr = sg.nr*vg.nr;
			//uint wnt = sg.nr*vg.nt;
			
			
			double *wr = new double[wnr];
			//double *wt = new double[wnt];
			
			memset(wr, 0, wnr*sizeof(*wr) );
			//memset(wt, 0, wnt*sizeof(*wt) );
			
			
			
			for (int i=0; i<n; i++) {
				
				double  r = sqrt(SQR(p[i].Pos[0])+SQR(p[i].Pos[1])+SQR(p[i].Pos[2])); 
				double _r = r*(sg.r2-sg.r1)/sg.r2 + sg.r1;
				uint ir = sg.nr*log(_r/sg.r1)/log(sg.r2/sg.r1);
			
				if (sg.nr <= ir) continue;
				
				Vector<REAL> V = *(Vector<REAL>*)p[i].Vel / sqrt(-2*p[i].pot);
				
				Vector<REAL> R(p[i].Pos[0],p[i].Pos[1], p[i].Pos[2]);
				//Vector<REAL> Z(0,0,1);
				//Vector<REAL> T(Z%R);
				
				double vr = V*R.ort();
				//double vt = V*T.ort();
				
				
				//double rc = (sg.r1*pow(sg.r2/sg.r1,(double)(i+0.5)/sg.nr)-sg.r1)*sg.r2/(sg.r2-sg.r1);
				//double v_max  = 1;
				
				double dvr = 2.0/vg.nr;
				//double dvt = 2.0/vg.nt;
				
				
				int vrid = int(floor(vr/dvr) ) + vg.nr/2;
				//int vtid = int(floor(vt/dvt) ) + vg.nt/2;
				
				
				if (0 <= vrid && vrid < vg.nr) wr[ir*vg.nr+vrid]++;
				//if (0 <= vtid && vtid < vg.nt) wt[ir*vg.nt+vtid]++;
				
				
			}
			
			for (int i=0; i<sg.nr; i++){
				double wrs(0), wts(0), wzs(0); 
				for (int k=0; k<vg.nr; k++) wrs+= wr[i*vg.nr+k];
				//for (int k=0; k<vg.nt; k++) wts+= wt[i*vg.nt+k];
				
				
				for (int k=0; k<vg.nr; k++) wr[i*vg.nr+k]/=wrs;
				//for (int k=0; k<vg.nt; k++) wt[i*vg.nt+k]/=wts;
				
			}
			
			
			mkdir((string(fnr).substr(0, string(fnr).find_last_of("/"))).c_str(), 0777);
			//mkdir((string(fnt).substr(0, string(fnt).find_last_of("/"))).c_str(), 0777);
			//mkdir((string(fnz).substr(0, string(fnz).find_last_of("/"))).c_str(), 0777);

			
			// saving
			FILE *fpr = fopen(fnr, "w");
			//FILE *fpt = fopen(fnt, "w");
			//FILE *fpz = fopen(fnz, "w");
			
			for (int i=0; i<sg.nr; i++) {
				for (int k=0; k<vg.nr; k++) fprintf(fpr, "%u %g\n", i*vg.nr+k, wr[i*vg.nr+k]);
				//for (int k=0; k<vg.nt; k++) fprintf(fpt, "%u %g\n", i*vg.nt+k, wt[i*vg.nt+k]);
				//for (int k=0; k<vg.nz; k++) fprintf(fpz, "%u %g\n", i*vg.nz+k, wz[i*vg.nz+k]);
			}
			
			
			fclose(fpr);
			//fclose(fpt);
			//fclose(fpz);
			
			
			delete [] wr;
			//delete [] wt;
			//delete [] wz;
			
			printf("done\n");
			
			

		}

		
		void calc_and_save_dm(const char *filename){
			
			printf("Halo::calc_and_save_reldendev('%s')...", filename); fflush(stdout);

			double r1 = 1;
			double r2 = 8192;
			uint   nr = 50;
			
			double *m0 = new double[nr];
			double *ms = new double[nr];
			uint *ns = new uint[nr];
			
			memset(ms, 0, nr*sizeof(double));
			memset(m0, 0, nr*sizeof(double));
			memset(ns, 0, nr*sizeof(uint));
		
			
			for (int i=0; i<n; i++) {
				
				double  r = sqrt(SQR(p[i].Pos[0])+SQR(p[i].Pos[1])+SQR(p[i].Pos[2])); 
				double _r = r*(r2-r1)/r2 + r1;
				uint ir = nr*log(_r/r1)/log(r2/r1);
				
				if (ir<nr) ms[ir]+=m/n;
				if (ir<nr) ns[ir]++;
				
			}
			double mp(0);
			double ma(0);
			FILE *fp = fopen(filename, "w");
			for (int i=0; i<nr; i++) {
				
				double _r  = (r1*pow(r2/r1,(double)(i+0.0)/nr)-r1)*r2/(r2-r1);
				double _r_ = (r1*pow(r2/r1,(double)(i+0.5)/nr)-r1)*r2/(r2-r1);
				double  r_ = (r1*pow(r2/r1,(double)(i+1.0)/nr)-r1)*r2/(r2-r1);
				
				double mi = M(_r, r_, 100);
				
				fprintf(fp, "%i %g %g %g %g %i\n",  i, _r_, ms[i], mi, (ms[i]-mi)/ms[i],  ns[i]);
				ma+=mi;
				mp+=ms[i];
			}
			
			fclose(fp);
			
			delete [] m0;
			delete [] ms;
			delete [] ns;
			
			printf("done\n");

			print();
			
			printf("m  = %.15f\n", m);
			printf("ma = %.15f\n", ma);
			printf("ms = %.15f\n", mp);

			
		}

		void center_potmin() {
			
			TreeForceFinder tff;
			tff.G = G;
			tff.NumPart = n;
			tff.force_treeallocate(1.5*n, n);
			tff.force_treebuild(p);
			
			int ci; 
			REAL acc[3];
			double minpot(1e99);
			for(int i=0; i<n; i++) {
				tff.force_treeevaluate(p, p[i].Pos, acc, p[i].pot);
				if (p[i].pot<minpot) {minpot = p[i].pot; ci = i; } 
			}
			tff.force_treefree();
			
			for(int i=0; i<n; i++) 
				*(Vector<REAL>*)p[i].Pos-=*(Vector<REAL>*)p[ci].Pos; 
				
		}

		void center_cm() {
			
					
			Vector<REAL> CM; double M(0);
			
			for (uint i=0; i<n; i++) {
		
				CM+=p[i].Mass * (*(Vector<REAL>*)p[i].Pos);
				M +=p[i].Mass;

			}
			
			CM/=M;
			CM.print();
			
			for(int i=0; i<n; i++) 
				*(Vector<REAL>*)p[i].Pos-=CM; 
				
		}
		
		void set_vel_cirlce(P *p, int n) {
			
			/*
			printf("Halo::set_vel_cirlce()... "); fflush(stdout);
			
			int i(0);
			while ( i<n) {
				
				double r = sqrt(SQR(p[i].Pos[0]) + SQR(p[i].Pos[1]) + SQR(p[i].Pos[2]));
				
				double v_cir = sqrt(-pot(r));
				
				p[i].Vel[0] = v_esc*(2*drand48() - 1);
				p[i].Vel[1] = v_esc*(2*drand48() - 1);
				p[i].Vel[2] = v_esc*(2*drand48() - 1);
				
				double v = sqrt(SQR(p[i].Vel[0]) + SQR(p[i].Vel[1]) + SQR(p[i].Vel[2]));
				if ( v_esc < v ) continue; 
				
				i++;
			}
			printf("done\n");
			*/
		}

		void remove_drift(){
			
			Vector<REAL> V;
			for (int i=0; i<n; i++)
				V+= *(Vector<REAL>*)p[i].Vel;
			
			V/=n;
			
			for (int i=0; i<n; i++)
				*(Vector<REAL>*)p[i].Vel -=V;
		
			V.print();
		}
		
		void save_pos(const char *filename) {
			printf("Disk::save_pos_float('%s')...", filename); fflush(stdout);
			FILE *fp = fopen(filename, "wb");
			for (int i=0; i<n; i++) {
				float pos[] = {p[i].Pos[0], p[i].Pos[1], p[i].Pos[2], 1 };
				fwrite(pos, sizeof(float), 4, fp);	
			}
			fclose(fp);
			printf("done\n");
		}

		void save_vel(const char *filename) {
			printf("Halo::save_vel_float('%s')...", filename); fflush(stdout);
			FILE *fp = fopen(filename, "wb");
			for (int i=0; i<n; i++) {
				float vel[] = {p[i].Vel[0], p[i].Vel[1], p[i].Vel[2], 1 };
				fwrite(vel, sizeof(float), 4, fp);	
			}
			fclose(fp);
			printf("done\n");
		}
		
		void load_vel(const char *filename) {
			
			printf("Halo::load_vel_float('%s')...", filename); fflush(stdout);
			FILE *fp = fopen(filename, "rb");
			
			for (int i=0; i<n; i++) {
				
				float vel[4];
				fread(vel, sizeof(float), 4, fp);
				
				p[i].Vel[0] = vel[0];
				p[i].Vel[1] = vel[1];
				p[i].Vel[2] = vel[2];
				
			}
			fclose(fp);
			printf("done\n");
		}

		void load_pos(const char *filename) {
			printf("Halo::load_pos_float('%s')...", filename); fflush(stdout);
			FILE *fp = fopen(filename, "rb");
			for (int i=0; i<n; i++) {
				
				float pos[4];
				fread(pos, sizeof(float), 4, fp);
				
				p[i].Pos[0] = pos[0];
				p[i].Pos[1] = pos[1];
				p[i].Pos[2] = pos[2];
				
			}
			fclose(fp);
			printf("done\n");
		}

		
		
		~Halo () {
			free();
		}

		void free() {
			if (p) delete [] p;
			p = NULL;
		}


		void print(){
			printf("Halo::print()\n", m);
			printf(" ├ m = %g\n", m);
			printf(" ├ r0 = %g\n", r0);
			printf(" └ rmax = %g\n\n", rmax);

		}
		
		
};

#endif
